import torch
def select_random_rows(tensor, num_rows):
    """
    从 tensor 中随机选择 num_rows 行进行训练。

    参数:
    tensor (torch.Tensor): 输入张量，维度为 [batch_size, n_vertices, ...]。
    num_rows (int): 要选择的行数。

    返回:
    torch.Tensor: 随机选择的行，维度为 [batch_size, num_rows, ...]。
    """
    n_vertices = tensor.size(1)
    
    # 生成不重复的随机索引
    indices = torch.randperm(n_vertices, device=tensor.device)[:num_rows]
    
    # 使用 torch.index_select 选择行
    return tensor.index_select(1, indices)

def select_random_rows_for_two_tensors(tensor1, tensor2, num_rows):
    n_vertices = tensor1.size(1)
    
    # 生成不重复的随机索引
    indices = torch.randperm(n_vertices, device=tensor1.device)[:num_rows]
    
    # 使用 torch.index_select 选择行
    return tensor1.index_select(1, indices),tensor2.index_select(1, indices)